{-# LANGUAGE TypeFamilies, TemplateHaskell, ViewPatterns, TypeOperators, Rank2Types #-}

{- |

Module      :  Type.Yoko.SetModelTH
Copyright   :  (c) The University of Kansas 2011
License     :  BSD3

Maintainer  :  nicolas.frisby@gmail.com
Stability   :  experimental
Portability :  see LANGUAGE pragmas (... GHC)

Some sets have obvious ':::' instances; this TH function derives them
automatically.

-}

module Type.Yoko.SetModelTH where

import Language.Haskell.TH

import Type.Booleans
import Type.Spine.TH (tvb_kind)

import Data.Maybe (catMaybes)
import Data.List (nub)
import Control.Arrow ((&&&))
import Control.Monad (liftM2)

import Type.Yoko.Type
import Type.Yoko.SetModel
import Type.Yoko.Wrapper

withTyVar :: String -> (Q Type -> a) -> a
withTyVar s = ($ newName s >>= varT)



set_declaration_information :: (forall a. Q a) -> Name -> Q (Q Type, [Name], [Q [(Type, Type)]], [Pred], Name, [Q Exp])
set_declaration_information bad n = do
  (tvbs, con) <- reify n >>= \i -> case i of
    TyConI dec -> case dec of
      DataD _ _ tvbs [con] _ -> return (tvbs, con)
      NewtypeD _ _ tvbs con _ -> return (tvbs, con)
      _ -> bad
    _ -> bad
  case tvbs of
    [] -> bad
    ((tvb_kind . last) -> StarK) -> return ()
    _ -> bad
  let ns = map tvb_name tvbs where
        tvb_name (PlainTV n) = n; tvb_name (KindedTV n _) = n
      tss = foldl (\t n -> t `appT` varT n) (conT n) (init ns)
      (cxt, ctor, argsRaw) = w con where
        w (NormalC n args) = ([], n, map snd args)
        w (RecC n args) = ([], n, map (\(_, _, z) -> z) args)
        w (InfixC argL n argR) = ([], n, map snd [argL, argR])
        w (ForallC _ cxt con) = (cxt ++ cxt', n, args)
          where (cxt', n, args) = w con

      args = map each argsRaw where
        -- TODO elaborate all type synonyms instead of just SuchThat
        each (AppT (AppT (ConT n) u) p) | n == ''SuchThat =
          AppT (AppT (ConT ''(,)) (AppT u wit)) (AppT p wit)
          where wit = AppT (AppT (ConT ''Witness) u) p
        each t = t

      (subUs, fields') = (map fst &&& map snd) $ map w args where
        w (AppT (AppT (ConT n) x) y) | n == ''(,) =
          (liftM2 (++) c1 c2, [|($fds1, $fds2)|])
          where (c1, fds1) = w x; (c2, fds2) = w y
        w (AppT u t) = (return $ [(u, t)], [|membership|])
        w _ = (bad, bad)
  return (tss, ns, subUs, cxt, ctor, fields')



obvious_membership = obvious_membership_ True True

obvious_membership_ :: Bool -> Bool -> Name -> Q [Dec]
obvious_membership_ genPred genMembership n = do
  let bad :: Q a
      bad = fail "`obvious_membership' expects the name of a data with exactly one constructor or newtype with at least one argument and its last argument is a *"

  (tss, ns, subUs, cxt, ctor, fields) <- set_declaration_information bad n

  subUs <- (nub . concat) `fmap` sequence subUs
  let cxt' = map (\(u, t) -> ClassP ''(:::) [t, u]) $ subUs
  tss <- return `fmap` tss
  uT <- return `fmap` varT (last ns)

  d <- fmap (:[]) $ 'membership `__varD` foldl appE (conE ctor) fields
  sequence $ catMaybes [
    if genMembership then Just $
      instanceD (return $ (nub . (cxt ++)) cxt')
        (conT ''(:::) `appT` uT `appT` tss) (map return d)
    else Nothing,
    if genPred then Just $
       tySynInstD ''(::?) [uT, tss] $
         let snoc sofar (u, t) = [t|And $sofar ($(return t) ::? $(return u))|]
         in foldl snoc (conT ''True) subUs
    else Nothing]

finite_set = finite_set_except []

finite_set_except :: [Name] -> Q Type -> Q [Dec]
finite_set_except except ty = do
  let dName :: Dec -> Name
      dName (TySynInstD n _ _) = n
      dName (InstanceD _ ty _) = w ty where
        w (ConT n) = n;   w (AppT ty _) = w ty

  tss <- return `fmap` ty

  t <- return `fmap` (newName "t" >>= varT)
  v <- return `fmap` (newName "v" >>= varT)

  pred0 <- classP ''(:::) [t, [t|Fin $tss|]]
  ds0 <- fmap ((:[]) . return) $ 'membership `__varD` appE (varE 'frFin) (varE 'membership)

  ds1 <- fmap ((:[]) . return) $ 'eqT `__varD` varE 'eqTFin

  pred2 <- classP ''(:=>) [[t|Fin $tss|], v]
  ds2 <- fmap ((:[]) . return) $ 'weaken `__varD` varE 'weakenFin

  ds3 <- fmap ((:[]) . return) $ 'evidence `__varD` [|\u x -> evidence (toFin u) x|]

  filter ((`notElem` except) . dName) `fmap` sequence [
    instanceD (return [pred0]) (appT (appT (conT ''(:::)) t) tss) ds0,
    instanceD (return []) (appT (conT ''EqT) tss) ds1,
    instanceD (return [pred2]) (appT (appT (conT ''(:=>)) tss) v) ds2,
    instanceD (return []) (appT (conT ''SetModel) tss) ds3,
    tySynInstD ''WitnessBT [tss, v] [t|WitnessBTFin $tss $v|],
    tySynInstD ''(::?) [t, tss] [t|$t ::? Fin $tss|]
           ]
